from flask import Flask, render_template, request, jsonify
import base64
import requests
import json

app = Flask(__name__)

# 百度AI的AK和SK
AK = 'YVkOOAXVUYQrU9eTjKjuFs5v'
SK = 'fhCGekZsQxng6tpHsLN9VoQ9a44qmAYY'

# 获取token
def get_access_token():
    url = 'https://aip.baidubce.com/oauth/2.0/token'
    params = {
        'grant_type': 'client_credentials',
        'client_id': AK,
        'client_secret': SK,
    }
    response = requests.get(url, params=params)
    return response.json().get('access_token')

# 文字识别
def ocr_text(image_base64, token):
    url = 'https://aip.baidubce.com/rest/2.0/ocr/v1/general_basic'
    params = {
        'access_token': token
    }
    data = {
        'image': image_base64
    }
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded'
    }
    response = requests.post(url, params=params, data=data, headers=headers)
    return response.json()

# 图片识别
def image_classify(image_base64, token):
    url = 'https://aip.baidubce.com/rest/2.0/image-classify/v2/advanced_general'
    params = {
        'access_token': token
    }
    data = {
        'image': image_base64
    }
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded'
    }
    response = requests.post(url, params=params, data=data, headers=headers)
    return response.json()

@app.route('/')
def index():
    return '''
    <!DOCTYPE html>
    <html lang="en">

    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>AI识别</title>
        <style>
            body {
                font-family: Arial, sans-serif;
                background-color: #f4f4f4;
                margin: 0;
                padding: 0;
                display: flex;
                justify-content: center;
                align-items: center;
                height: 100vh;
                flex-direction: column;
            }
            h1 {
                color: #333;
            }
            form {
                background: #fff;
                padding: 20px;
                border-radius: 8px;
                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
                width: 300px;
                text-align: center;
            }
            input[type="file"], input[type="checkbox"] {
                margin: 10px 0;
            }
            label {
                display: block;
                margin: 10px 0;
            }
            button {
                background-color: #28a745;
                color: white;
                padding: 10px 20px;
                border: none;
                border-radius: 4px;
                cursor: pointer;
            }
            button:hover {
                background-color: #218838;
            }
            .result {
                margin-top: 20px;
                background: #fff;
                padding: 20px;
                border-radius: 8px;
                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
                width: 300px;
                text-align: left;
            }
            .result h2 {
                margin-top: 0;
                color: #333;
            }
            .result p {
                margin: 5px 0;
            }
        </style>
    </head>

    <body>
        <h1>上传图片进行识别</h1>
        <form action="/upload" method="post" enctype="multipart/form-data">
            <input type="file" name="file" accept="image/*">
            <label>
                <input type="checkbox" name="ocr" value="1"> 文字识别
            </label>
            <label>
                <input type="checkbox" name="classify" value="1"> 图片分类
            </label>
            <button type="submit">上传</button>
        </form>
    </body>

    </html>
    '''

@app.route('/upload', methods=['POST'])
def upload():
    file = request.files['file']
    image = file.read()
    image_base64 = base64.b64encode(image).decode('utf-8')
    
    token = get_access_token()
    
    ocr_result = None
    classify_result = None
    
    if 'ocr' in request.form:
        ocr_result = ocr_text(image_base64, token)
    if 'classify' in request.form:
        classify_result = image_classify(image_base64, token)
    
    # 提取 OCR 结果中的文字
    ocr_words = [word['words'] for word in ocr_result.get('words_result', [])] if ocr_result else []
    
    # 提取图像分类结果中的标签和置信度
    classify_labels = [f"{item['keyword']} (置信度: {item['score']:.2f})" for item in classify_result.get('result', [])] if classify_result else []
    
    return f'''
    <!DOCTYPE html>
    <html lang="en">

    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>AI识别结果</title>
    </head>

    <body>
        <div class="result">
            <h2>文字识别结果</h2>
            {"<ul><li>" + "</li><li>".join(ocr_words) + "</li></ul>" if ocr_words else "未识别到文字"}
        </div>
        <div class="result">
            <h2>图像分类结果</h2>
            {"<ul><li>" + "</li><li>".join(classify_labels) + "</li></ul>" if classify_labels else "未识别到图像类别"}
        </div>
    </body>

    </html>
    '''

if __name__ == '__main__':
    app.run(debug=True)